# main.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import gradio as gr
import numpy as np
from PIL import Image

# 定义MLP模型结构
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten the input
        x = F.relu(self.fc1(x))
        x = F.dropout(x, p=0.2, training=self.training)
        x = F.relu(self.fc2(x))
        x = F.dropout(x, p=0.2, training=self.training)
        x = self.fc3(x)
        return F.log_softmax(x, dim=1)

# 加载MNIST数据集并进行预处理
"""transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Normalize((0.1307,), (0.3081,))])

train_set = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=True, transform=transform)
test_set = datasets.MNIST('~/.pytorch/MNIST_data/', download=True, train=False, transform=transform)

x_train, y_train = train_set.data, train_set.targets
x_test, y_test = test_set.data, test_set.targets

# 转换数据类型
x_train, x_test = x_train.float(), x_test.float()

# 构建MLP模型
model = MLP()

# 设置损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
model.train()  # Set the model to training mode
for epoch in range(10):
    optimizer.zero_grad()
    output = model(x_train)
    loss = criterion(output, y_train)
    loss.backward()
    optimizer.step()

# 保存模型参数
torch.save(model.state_dict(), 'best_mlp.pth')"""

# 定义预处理函数
def preprocess(image):
    image = Image.fromarray(image).convert('L')  # 转换为灰度图像
    image_tensor = transforms.ToTensor()(image)  # 转换为张量
    image_tensor = image_tensor.view(1, -1)  # 添加批次维度
    return image_tensor

# 定义预测函数
def predict(image):
    preprocessed_image = preprocess(image)
    model.eval()  # Set the model to evaluation mode
    with torch.no_grad():
        output = model(preprocessed_image)
        _, predicted_digit = torch.max(output, 1)
    return predicted_digit.item()

# 创建Gradio接口
iface = gr.Interface(fn=predict, inputs='sketchpad', outputs='text')

# 启动Gradio接口
iface.launch()
